Skip to content

Comments

Optimize Kafka IO: Remove Guava Stopwatch overhead from consumer.poll() loop#37705

Open
junaiddshaukat wants to merge 1 commit intoapache:masterfrom
junaiddshaukat:remove-kafka-stopwatch-overhead
Open

Optimize Kafka IO: Remove Guava Stopwatch overhead from consumer.poll() loop#37705
junaiddshaukat wants to merge 1 commit intoapache:masterfrom
junaiddshaukat:remove-kafka-stopwatch-overhead

Conversation

@junaiddshaukat
Copy link
Contributor

@junaiddshaukat junaiddshaukat commented Feb 25, 2026

Fixes #37704

What happened?

This PR optimizes the Kafka IO readers (ReadFromKafkaDoFn and KafkaUnboundedReader) by removing the Guava Stopwatch and the custom RPC latency metric tracking from the hot consumer.poll() loop.

The Problem:
Guava's Stopwatch relies on System.nanoTime(), which is a relatively expensive system call. In the Kafka IO module, this timer was being started and stopped repeatedly inside the consumer.poll() loop. When Kafka prefetches records, this loop executes extremely rapidly, causing the Stopwatch to introduce significant and unnecessary CPU overhead.

The Solution:

  1. Removed Stopwatch: Replaced the Stopwatch timeout tracking with a lightweight System.currentTimeMillis() calculation to track the remainingTimeout.
  2. Removed Custom RPC Metrics: Removed the updateSuccessfulRpcMetrics calls and associated histograms (KafkaMetrics, KafkaSinkMetrics).
  3. Why this is safe: The custom Beam RPC latency metric is redundant. Kafka already natively tracks and exposes this data via JMX metrics (specifically fetch-latency-avg), so we can safely drop the custom implementation to regain performance.

Changes made:

  • Updated ReadFromKafkaDoFn.java and KafkaUnboundedReader.java to use System.currentTimeMillis() for timeout tracking instead of Stopwatch.
  • Cleaned up dead code in KafkaMetrics.java and KafkaSinkMetrics.java by removing RPC_LATENCY histograms and updateSuccessfulRpcMetrics.
  • Updated KafkaMetricsTest.java and KafkaSinkMetricsTest.java to reflect the removed metrics.

Testing Done

  • Compiled successfully.
  • Ran and passed the full Kafka IO test suite (./gradlew :sdks:java:io:kafka:test).
  • Verified that timeout logic in the poll loop still functions correctly using java.time.Duration and System.currentTimeMillis().

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

  • Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
  • Update CHANGES.md with noteworthy changes.
  • If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

Build python source distribution and wheels
Python tests
Java tests
Go tests

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

@gemini-code-assist
Copy link
Contributor

gemini-code-assist bot commented Feb 25, 2026

Summary of Changes

Hello @junaiddshaukat, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the performance of Kafka IO readers by addressing and removing two sources of unnecessary CPU overhead. By replacing an expensive timer mechanism with a more efficient alternative and eliminating redundant custom RPC latency metrics, the Kafka consumer's poll() loop now executes with reduced latency, especially when processing pre-fetched records. This change streamlines the data ingestion process from Kafka, leading to more efficient resource utilization.

Highlights

  • Performance Optimization in Kafka IO: Removed the computationally expensive Guava Stopwatch from the critical consumer.poll() loop in Kafka IO readers (ReadFromKafkaDoFn and KafkaUnboundedReader).
  • Replaced Stopwatch with System.currentTimeMillis(): Switched from Stopwatch (which relies on System.nanoTime()) to a more lightweight System.currentTimeMillis() for tracking remainingTimeout, significantly reducing CPU overhead during rapid Kafka record prefetching.
  • Removed Redundant RPC Latency Metrics: Eliminated custom Beam RPC latency metric tracking and associated histograms (KafkaMetrics, KafkaSinkMetrics) as Kafka natively provides this data via JMX (fetch-latency-avg), making the custom implementation redundant and a source of performance overhead.
  • Cleaned Up Metric Code and Tests: Removed dead code related to RPC latency metrics from KafkaMetrics.java and KafkaSinkMetrics.java, and updated corresponding test files (KafkaMetricsTest.java, KafkaSinkMetricsTest.java) to reflect these changes.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaMetrics.java
    • Removed updateSuccessfulRpcMetrics method signature from the interface.
    • Removed updateSuccessfulRpcMetrics implementation from NoOpKafkaMetrics.
    • Removed LATENCY_HISTOGRAMS map and perTopicRpcLatencies abstract method.
    • Adjusted create method in KafkaMetricsImpl to no longer initialize perTopicRpcLatencies.
    • Removed the updateSuccessfulRpcMetrics method implementation from KafkaMetricsImpl.
    • Removed the recordRpcLatencyMetrics method.
    • Removed the call to recordRpcLatencyMetrics from flushBufferedMetrics.
  • sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaSinkMetrics.java
    • Removed imports for DelegatingHistogram, Histogram, and HistogramData.
    • Removed RPC_LATENCY constant.
    • Removed RpcMethod enum.
    • Removed RPC_METHOD metric label.
    • Removed the createRPCLatencyHistogram method.
  • sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaUnboundedReader.java
    • Removed import for Stopwatch.
    • Removed the stopwatch instance variable.
    • Removed stopwatch.start(), stopwatch.stop(), and kafkaResults.updateSuccessfulRpcMetrics calls from the consumerPollLoop method.
  • sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/ReadFromKafkaDoFn.java
    • Removed import for Stopwatch.
    • Removed pollTimer (Stopwatch) initialization.
    • Replaced pollTimer.reset().start() with System.currentTimeMillis() for startMillis.
    • Replaced pollTimer.elapsed() with Duration.ofMillis(System.currentTimeMillis() - startMillis) for elapsed calculation.
    • Removed kafkaMetrics.updateSuccessfulRpcMetrics call.
  • sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaMetricsTest.java
    • Removed import for containsInAnyOrder.
    • Removed import for Duration.
    • Removed results.updateSuccessfulRpcMetrics call from testNoOpKafkaMetrics.
    • Renamed testKafkaRPCLatencyMetrics to testKafkaBacklogMetrics.
    • Removed all assertions related to RPC latency histograms from testKafkaBacklogMetrics.
    • Removed the testKafkaRPCLatencyMetricsAreNotRecorded test method.
  • sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaSinkMetricsTest.java
    • Removed import for Histogram.
    • Removed the testCreatingHistogram method.
Activity
  • The code compiled successfully.
  • The full Kafka IO test suite was run and passed.
  • The timeout logic in the poll loop was verified to function correctly using java.time.Duration and System.currentTimeMillis().
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

@github-actions
Copy link
Contributor

github-actions bot commented Feb 25, 2026

Checks are failing. Will not request review until checks are succeeding. If you'd like to override that behavior, comment assign set of reviewers

Abacn
Abacn requested changes Feb 25, 2026
View reviewed changes
Copy link
Contributor

@Abacn Abacn left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's deleting functioning metrics and codes. Please check generated code changes before pushing PRs

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@Abacn Abacn Abacn requested changes

Assignees

No one assigned

Labels

io java kafka

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Task]: Kafka IO: Remove custom latency timer in ReadFromKafkaDoFn in favor Kafka's native metrics

2 participants

@junaiddshaukat @Abacn